package tree.huffman.code;

import java.io.*;
import java.util.*;

/**
 * @ author : HyoJoo-W
 * @ createDate : 2021-06-05 15:40
 * @ Description :
 */
public class HuffmanCode
{
    public static void main(String[] args)
    {
        ////初始
        //String str = "i like like like java do you like a java";
        //byte[] bytes = str.getBytes();
        //System.out.println("压缩前:" + Arrays.toString(bytes));
        //System.out.println("压缩前:" + bytes.length);
        //
        ////生成哈夫曼树
        //List<Node> nodes = new ArrayList<>();
        //nodes = getNodes(bytes);
        //Node root = createHuffmanTree(nodes);
        ////preOrder(root);
        //
        ////生成哈夫曼编码
        //Map<Byte, String> huffmanCodes = getHuffmanCodes(root);
        ////System.out.println("Huffman编码表:" + huffmanCodes);
        //
        ////使用哈夫曼编码对原字符串压缩
        //byte[] compress = compress(bytes, huffmanCodes);
        //System.out.println("压缩前:" + Arrays.toString(compress));
        //System.out.println("压缩前:" + compress.length);
        //
        //float compressRate = (float) (bytes.length - compress.length) / bytes.length;
        //System.out.println("压缩率: " + compressRate);
        //
        ////String strByte = "10101000";
        ////System.out.println(Integer.parseInt(strByte,2));
        ////System.out.println(((byte) Integer.parseInt(strByte, 2)));
        //
        //System.out.println("解压后...");
        //System.out.println(new String(decode(huffmanCodes, compress)));
        ////System.out.println(decode(huffmanCodes, compress).length);

        //String src = "E://test.txt";
        //String target = "E://test.huf";
        //
        //HuffmanCompressFile(src, target);

        String src = "E://test.huf";
        String target = "E://decode.txt";
        HuffmanUnCompressFile(src,target);


    }

    /**
     * @param bytes: 原始字符串对应的byte数组
     * @return byte[]
     * @Description 总方法
     **/
    public static byte[] huffmanCompress(byte[] bytes)
    {
        List<Node> nodes = getNodes(bytes);
        Node root = createHuffmanTree(nodes);
        Map<Byte, String> huffmanCodes = getHuffmanCodes(root);

        return compress(bytes, huffmanCodes);
    }

    /**
     * @param flag: 是否需要补位(正数)
     * @return 转换后的二进制字符串
     * @Description 字节转二进制字符串
     **/
    private static String byteToBinaryString(boolean flag, byte b)
    {
        int temp = b;
        if (flag)
        {
            temp |= 256;
        }
        String str = Integer.toBinaryString(temp);
        if (flag)
        {
            return str.substring(str.length() - 8);
        }
        else
        {
            return str;
        }
    }

    //解码
    private static byte[] decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes)
    {
        //1. 得到huffmanBytes对应的二进制字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++)
        {
            //判断是否是最后一个字节
            boolean flag = (i == huffmanBytes.length - 1);
            stringBuilder.append(byteToBinaryString(!flag, huffmanBytes[i]));
        }

        //2. 将huffmanCode的map反转
        Map<String, Byte> reverseMap = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet())
        {
            reverseMap.put(entry.getValue(), entry.getKey());
        }

        List<Byte> bytes = new ArrayList<>();
        for (int i = 0; i < stringBuilder.length(); )
        {
            int count = 1;
            boolean flag = true;
            Byte b = null;

            while (flag)
            {
                String key = stringBuilder.substring(i, i + count);
                b = reverseMap.get(key);
                if (b == null)
                {
                    count++;
                }
                else
                {
                    flag = false;
                }
            }
            bytes.add(b);
            i += count;//控制步长
        }

        //返回
        byte[] returnBytes = new byte[bytes.size()];
        for (int i = 0; i < returnBytes.length; i++)
        {
            returnBytes[i] = bytes.get(i);
        }
        return returnBytes;

    }

    private static void preOrder(Node root)
    {
        if (root != null)
        {
            root.preOrder();
        }
        else
        {
            System.out.println("根节点为空");
        }
    }

    //根据待编码的字节数组生成节点
    private static List<Node> getNodes(byte[] bytes)
    {
        List<Node> nodes = new ArrayList<>();

        Map<Byte, Integer> weights = new HashMap<>();

        for (byte aByte : bytes)
        {
            Integer count = weights.get(aByte);
            if (count == null)
            {
                weights.put(aByte, 1);
            }
            else
            {
                weights.put(aByte, count + 1);
            }
        }

        //把map中的每一对键值对转成一个Node对象,加入到nodes中
        for (Map.Entry<Byte, Integer> entry : weights.entrySet())
        {
            nodes.add(new Node(entry.getKey(), entry.getValue()));
        }

        return nodes;
    }

    //根据节点生成Huffman树
    private static Node createHuffmanTree(List<Node> nodes)
    {
        while (nodes.size() > 1)
        {
            Collections.sort(nodes);

            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            //新二叉树的根节点没有data!
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;

            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    //存放Huffman编码
    static Map<Byte, String> huffmanCodes = new HashMap<>();
    //存储某个节点的路径
    static StringBuilder stringBuilder = new StringBuilder();

    /**
     * @param node:          传入节点
     * @param code:          路径对应的编码,向左为0,向右为1
     * @param stringBuilder: 拼接路径
     * @return void
     * @Description 求出将传入的node节点的所有叶子节点的哈夫曼编码, 并放入huffmanCodes
     **/
    private static void getHuffmanCodes(Node node, String code, StringBuilder stringBuilder)
    {
        StringBuilder stringBuilder2 = new StringBuilder(stringBuilder);
        stringBuilder2.append(code);
        if (node != null)
        {
            //判断node是否为叶子节点
            if (node.data == null)
            {
                //向左
                getHuffmanCodes(node.left, "0", stringBuilder2);
                //向右
                getHuffmanCodes(node.right, "1", stringBuilder2);
            }
            else
            {
                //递归终止条件: 遍历到叶子节点时, 不在继续遍历,当前层数的程序执行结束,返回上一层
                //叶子节点
                huffmanCodes.put(node.data, stringBuilder2.toString());
            }
        }
    }

    private static Map<Byte, String> getHuffmanCodes(Node root)
    {
        if (root == null)
        {
            return null;
        }
        //左子树
        getHuffmanCodes(root.left, "0", stringBuilder);
        //右子树
        getHuffmanCodes(root.right, "1", stringBuilder);

        return huffmanCodes;
    }


    /**
     * @param bytes:        原始字符串的byte数组
     * @param huffmanCodes: 原始字符串生成的哈夫曼编码表
     * @return 根据哈夫曼编码后的字符串为1、0串,对该串(补码)求原码得到byte数组
     * @Description 压缩
     **/
    private static byte[] compress(byte[] bytes, Map<Byte, String> huffmanCodes)
    {
        StringBuilder stringBuilder = new StringBuilder();
        for (byte aByte : bytes)
        {
            stringBuilder.append(huffmanCodes.get(aByte));
        }
        //System.out.println(stringBuilder);

        //统计返回的byte数组长度
        int length;
        if (stringBuilder.length() % 8 == 0)
        {
            length = stringBuilder.length() / 8;
        }
        else
        {
            length = stringBuilder.length() / 8 + 1;
        }

        //int length = (stringBuilder.length() + 7) / 8

        byte[] compressed = new byte[length];
        int byteIndex = 0;
        for (int i = 0; i < stringBuilder.length(); i += 8)
        {
            String strByte;
            if ((i + 8) > stringBuilder.length())
            {
                strByte = stringBuilder.substring(i);
            }
            else
            {
                strByte = stringBuilder.substring(i, i + 8);
            }
            compressed[byteIndex++] = ((byte) Integer.parseInt(strByte, 2));
        }
        return compressed;
    }


    /**
     * @param src:    源文件目录
     * @param target: 压缩后的文件目录
     * @return
     * @Description 文件压缩
     **/
    public static void HuffmanCompressFile(String src, String target)
    {
        try
        {
            FileInputStream fis = new FileInputStream(src);
            byte[] bytes = new byte[fis.available()];
            fis.read(bytes);

            byte[] huffmanBytes = huffmanCompress(bytes);

            FileOutputStream fos = new FileOutputStream(target);
            ObjectOutputStream oos = new ObjectOutputStream(fos);
            oos.writeObject(huffmanBytes);//把压缩后的字节数组以对象流写入压缩文件
            oos.writeObject(huffmanCodes);//把哈夫曼编码以对象流写入压缩文件(解压时使用)

            fis.close();
            fos.close();
            oos.close();
        }
        catch (IOException e)
        {
            e.printStackTrace();
        }
    }

    /**
     * @Description 解压
     * @param src: 源文件
     * @return void
     **/
    public static void HuffmanUnCompressFile(String src, String target)
    {
        InputStream is = null;
        ObjectInputStream ois = null;
        OutputStream os = null;
        try
        {
            is = new FileInputStream(src);
            ois = new ObjectInputStream(is);
            byte[] huffmanBytes = ((byte[]) ois.readObject());
            Map<Byte, String> huffmanCodes = ((Map<Byte,String>) ois.readObject());
            byte[] decode = decode(huffmanCodes, huffmanBytes);

            os = new FileOutputStream(target);
            os.write(decode);

            is.close();
            ois.close();
            os.close();
        }
        catch (IOException | ClassNotFoundException e)
        {
            e.printStackTrace();
        }
    }
}


class Node implements Comparable<Node>
{
    public Byte data;
    public int weight;
    public Node left;
    public Node right;

    public Node(Byte data, int weight)
    {
        this.data = data;
        this.weight = weight;
    }

    @Override
    public int compareTo(Node o)
    {
        return this.weight - o.weight;
    }

    @Override
    public String toString()
    {
        return "Node{" +
                "data=" + data +
                ", weight=" + weight +
                '}';
    }

    public void preOrder()
    {
        System.out.println(this);
        if (this.left != null)
        {
            this.left.preOrder();
        }
        if (this.right != null)
        {
            this.right.preOrder();
        }
    }


}

